from langchain_core.prompts import ChatPromptTemplate
from langchain_deepseek import ChatDeepSeek
from typing import Optional
from typing_extensions import Annotated, TypedDict
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

llm = ChatDeepSeek(
    model="deepseek-chat",  # 模型名称
    temperature=0,         # 控制生成随机性(0-1)
    max_tokens=None,       # 最大输出token数
    timeout=None,          # 超时设置
    max_retries=2,         # 失败重试次数
    api_key="sk-XXXXXXXX"
)

class Joke(TypedDict):
    """Joke to tell user."""

    setup: Annotated[str, ..., "The setup of the joke"]

    # Alternatively, we could have specified setup as:

    # setup: str                    # no default, no description
    # setup: Annotated[str, ...]    # no default, no description
    # setup: Annotated[str, "foo"]  # default, no description

    punchline: Annotated[str, ..., "The punchline of the joke"]
    rating: Annotated[Optional[int], None, "How funny the joke is, from 1 to 10"]
structured_llm = llm.with_structured_output(Joke)
# few shots
system = """You are a hilarious comedian. Your specialty is knock-knock jokes. \
Return a joke which has the setup (the response to "Who's there?") and the final punchline (the response to "<setup> who?").

Here are some examples of jokes:

example_user: Tell me a joke about planes
example_assistant: {{"setup": "Why don't planes ever get tired?", "punchline": "Because they have rest wings!", "rating": 2}}

example_user: Tell me another joke about planes
example_assistant: {{"setup": "Cargo", "punchline": "Cargo 'vroom vroom', but planes go 'zoom zoom'!", "rating": 10}}

example_user: Now about caterpillars
example_assistant: {{"setup": "Caterpillar", "punchline": "Caterpillar really slow, but watch me turn into a butterfly and steal the show!", "rating": 5}}"""
prompt = ChatPromptTemplate.from_messages([("system", system), ("human", "{input}")])
few_shot_structured_llm = prompt | structured_llm
print(few_shot_structured_llm.invoke({"input": "what's something funny about woodpeckers"}))